from typing import Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from gxl_ai_utils.utils import utils_file

from wenet.transformer.ctc import CTC


class CTCForLLM(nn.Module):
    def __init__(self, input_dim, output_dim, dropout=0.2, blank_id=0):
        super(CTCForLLM, self).__init__()
        self.ctc = CTC(odim=output_dim, encoder_output_size=input_dim, dropout_rate=dropout, reduce=True, blank_id=blank_id)

    def forward(self, hs_pad: torch.Tensor, hlens: torch.Tensor,
                ys_pad: torch.Tensor,
                ys_lens: torch.Tensor) -> torch.Tensor:
        """
        return loss value
        :param hs_pad: (B, Tmax, D)
        :param hlens: (B,)
        :param ys_pad: (B, Lmax)
        :param ys_lens: (B,)
        :return: loss value
        """
        return self.ctc(hs_pad, hlens, ys_pad, ys_lens)[0]


if __name__ == '__main__':
    model = CTCForLLM(3584, 152064)
    utils_file.print_model_size(model)